# -*- coding: utf-8 -*-
# @time: 2023/09/11 09:39
# @file: get_base_ret.py
# @author: tyshixi08


# 调取需要的模块
import pandas as pd
from datetime import date
import numpy as np
# import rqdatac
from sqlalchemy import create_engine
from tqdm import tqdm
import os
from WindPy import w
w.start()
# rqdatac.init('license', 'AcBHy5_JJ6wjZdu7Q-ey7dX-J3BmyEC_KblY2Q_hBeOuoBaeBbgXTNSe6XZvqKVESbyUf7vMpLLGuO_aqyb3w9fWGI7q4wdClE6cMp_Z3N4PqqTHJ0nr3CIuXtk-5XzSD1p7NTdNcrAfZlRVpMMtY_PDC9FYuXNmC_EnuQg4H-A=fGk9EhHcK3xN189iXYSWLyiMdGUeXXlVZqr2MxhBypSHxQYnIIyxyM8BR8oNnVUdWhKx-ZrFRIjSONd7uYpOvpcBab92P60iAR_JopX61emtrvsY1xG_uCfYhDPBdDSJKaniJhTPuoBIU4JZun8-8fMIxzx7lnwBm2kAUOA_Mpg=')


# def bond_convert(code):
#     """
#     从米筐合约代码转换至交易所合约代码(米筐仅支持转换A股、期货和期权)
#     """
#     if code[7:] == 'XSHG':
#         code = code[:7] + 'SH'
#     elif code[7:] == 'XSHE':
#         code = code[:7] + 'SZ'
#     return code


# %% 数据提取class ConvBondData
# 提取相关数据#
class ConvBondData(object):
    def __init__(self):
        #self.path = 'http://dataway.hhhstz.com:8888/hsic_base_fmt/cube?'
        self.path = 'http://dataway.hhhstz.com/hsic_base_fmt/cube?'
        # all_instruments = rqdatac.convertible.all_instruments()
        # all_instruments = all_instruments[['order_book_id', 'listed_date', 'maturity_date', 'stock_code']]
        # all_instruments = all_instruments.rename(columns={'order_book_id': 'c_bondCode', 'listed_date': 't_ipoDate',
        #                                                   'maturity_date': 't_maturityDate',
        #                                                   'stock_code': 'c_underlyingCode'})
        # self.all_instruments = all_instruments

    # 股票日行情估值表
    def get_stocka_derivativeindicator(self, start_date, end_date, fields=None):
        tableName = 'b_stocka_derivativeindicator'
        query_str = self.path + 'tableName=' + tableName + '&begDate=' + start_date + '&endDate=' + end_date
        if fields != None: query_str += '&fields=c_code,t_date,' + ','.join(fields)
        print(query_str)
        data = pd.read_csv(query_str)
        # engine = create_engine('mysql+pymysql://hs_wangwenjie:hs_wangwenjie#A0@192.168.201.181:3306/wind?charset=utf8')
        # query = ("""WindCustomCode""")
        # data = pd.read_sql_query(query, engine)
        return data

    # 股票日行情
    def get_stocka_marketday(self, start_date, end_date, fields=None):
        tableName = 'b_stocka_marketday'
        query_str = self.path + 'tableName=' + tableName + '&begDate=' + start_date + '&endDate=' + end_date
        if fields != None: query_str += '&fields=c_code,t_date,' + ','.join(fields)
        print(query_str)
        data = pd.read_csv(query_str)
        return data
    
    # 股票日行情估值指标表
    def get_stocka_plateinfo(self, start_date, end_date, fields=None):
        tableName = 'b_stocka_plateinfo'
        query_str = self.path + 'tableName=' + tableName + '&begDate=' + start_date + '&endDate=' + end_date
        if fields != None: query_str += '&fields=c_code,t_date,' + ','.join(fields)
        print(query_str)
        data = pd.read_csv(query_str)
        return data

    # 债券估值表 -> 米筐：convertible.all_instruments, convertible.get_indicators
    # 日期0，更新日期，转债代码0，标的代码0，(转债评级，发行人评级)，纯债溢价x，纯债溢价率8，纯债价值7，转股溢价x，转股溢价率3，转股价格x，
    # 转股比例1，转换价值2，当期收益率，转债余额5，正股市净率6，正股市盈率，纯债到期收益率4，更新日期，上市日期0，到期日期0，更新时间，隐含波动率9
    def get_convbond_valuation(self, start_date, end_date, fields=None):
        # # 基础指标
        # all_instruments = self.all_instruments
        # convbond_code = list(all_instruments.c_bondCode.unique())
        #
        # # 估值指标
        # indicators = rqdatac.convertible.get_indicators(convbond_code, start_date, end_date,
        #                                                 fields=['conversion_coefficient', 'conversion_value',
        #                                                         'conversion_premium', 'yield_to_maturity_pretax',
        #                                                         'remaining_size', 'pb_ratio', 'pure_bond_value_1',
        #                                                         'pure_bond_value_premium_1', 'iv'])
        # indicators = indicators.reset_index()
        # indicators = indicators.rename(
        #     columns={'order_book_id': 'c_bondCode', 'date': 'index', 'conversion_coefficient': 'n_convRatio',
        #              'conversion_value': 'n_convValue', 'conversion_premium': 'n_convPremiumRatio',
        #              'yield_to_maturity_pretax': 'n_ytmCB', 'remaining_size': 'n_outstandingBalance',
        #              'pb_ratio': 'n_underlyingPB', 'pure_bond_value_1': 'n_bondValue',
        #              'pure_bond_value_premium_1': 'n_bondPremiumRatio', 'iv': 'n_implied_vol'})
        # indicators['index'] = indicators['index'].apply(lambda x: x.strftime('%Y-%m-%d'))
        #
        # # 评级
        # # 读取已有评级数据
        # df_old_rating = pd.read_csv('./param/old_issuerRating2.csv')
        # df_old_rating = df_old_rating[(df_old_rating['index'] >= start_date) & (df_old_rating['index'] <= end_date)]
        # latest_day = df_old_rating['index'].unique()[-1]
        # # 从wind中获取最新评级数据
        # new_rating = w.wsd(convbond_code, "latestissurercreditrating2", latest_day, end_date, "ratingAgency = 101; type = 17")
        # df_new_rating = pd.DataFrame(sum(new_rating.Data, []),
        #                              index=pd.MultiIndex.from_product([new_rating.Codes, new_rating.Times]),
        #                              columns=['c_issuerRating2'])
        # df_new_rating.index.names = ['c_bondCode', 'index']
        # df_new_rating = df_new_rating.reset_index()[['index', 'c_bondCode', 'c_issuerRating2']]
        # df_new_rating['index'] = df_new_rating['index'].apply(lambda x: x.strftime('%Y-%m-%d'))
        # # 合并评级数据
        # df_rating = pd.concat([df_old_rating, df_new_rating]).sort_values('index').reset_index(drop=True)
        # df_rating.to_csv('./param/old_issuerRating2.csv')
        #
        # # 汇总数据
        # data = pd.merge(indicators, all_instruments, on='c_bondCode', how='left')
        # data['c_bondCode'] = data['c_bondCode'].apply(bond_convert)
        # data['c_underlyingCode'] = data['c_underlyingCode'].apply(bond_convert)
        # data = pd.merge(data, df_rating, on=['index', 'c_bondCode'], how='left')
        tableName = 'b_convbond_valuation'
        query_str = self.path + 'tableName=' + tableName + '&begDate=' + start_date + '&endDate=' + end_date
        if fields != None: query_str += '&fields=c_code,t_date,' + ','.join(fields)
        print(query_str)
        data = pd.read_csv(query_str)
        return data
    
    # 债券信息表
    def get_convbond_info(self, start_date, end_date, fields=None):
        tableName = 'b_bond_info'
        query_str = self.path + 'tableName=' + tableName + '&begDate=' + start_date + '&endDate=' + end_date
        if fields != None: query_str += '&fields=c_code,t_date,' + ','.join(fields)
        print(query_str)
        data = pd.read_csv(query_str)
        return data

    # 债券指数权重表 
    def get_index_weight(self, start_date, end_date, fields=None):
        tableName = 'b_convbond_indexweight'
        query_str = self.path + 'tableName=' + tableName + '&begDate=' + start_date + '&endDate=' + end_date
        if fields != None: query_str += '&fields=c_code,t_date,' + ','.join(fields)
        data = pd.read_csv(query_str)
        return data

    # 债券日行情表 -> 米筐: convertible.all_instruments, get_price
    # 日期1，转债代码1，利率，成交额1，涨跌cal，收盘价1，最高价1，最低价1，开盘价1，涨跌幅cal，前日收盘价1，成交量1
    def get_convbond_market(self, start_date, end_date, fields=None):
        # # 基础指标
        # all_instruments = self.all_instruments
        # convbond_code = list(all_instruments.c_bondCode.unique())
        # # 转债日行情
        # data = rqdatac.get_price(convbond_code, start_date=start_date, end_date=end_date, frequency='1d', fields=None,
        #                           adjust_type='none', skip_suspended=False, market='cn', expect_df=True, time_slice=None)
        # data = data[['close', 'high', 'low', 'open', 'total_turnover', 'volume', 'prev_close']]
        # data = data.reset_index()
        # data = data.rename(columns={'order_book_id': 'c_bondCode', 'date': 'index', 'close': 'n_close',
        #                             'high': 'n_high', 'low': 'n_low', 'open': 'n_open', 'total_turnover': 'n_amt',
        #                             'volume': 'n_volume', 'prev_close': 'n_preClose'})
        # data['n_change'] = data.groupby('c_bondCode').n_close.apply(lambda x: x - x.shift(1)).values
        # data['n_pctChange'] = data.groupby('c_bondCode').n_close.apply(lambda x: (x - x.shift(1)) / x.shift(1)).values
        # data['index'] = data['index'].apply(lambda x: x.strftime('%Y-%m-%d'))
        # data['c_bondCode'] = data['c_bondCode'].apply(bond_convert)
        tableName = 'b_convbond_marketday'
        query_str = self.path + 'tableName=' + tableName + '&begDate=' + start_date + '&endDate=' + end_date
        if fields != None: query_str += '&fields=c_code,t_date,' + ','.join(fields)
        data = pd.read_csv(query_str)
        return data

    # 正股停牌日 -> 米筐: convertible.all_instruments, is_suspended
    def get_stocka_warning(self, start_date, end_date, fields=None):
        # # 基础数据
        # all_instruments = self.all_instruments
        # stock_code0 = list(all_instruments['c_underlyingCode'].unique())
        # stock_code = [code if not pd.isnull(code) else None for code in stock_code0]
        # while None in stock_code:
        #     stock_code.remove(None)
        # # 正股停牌日
        # data = rqdatac.is_suspended(stock_code, start_date=start_date, end_date=end_date, market='cn')
        # data = data.stack().reset_index()
        # data.columns = ['index', 'c_code', 't_warning']
        # data['index'] = data['index'].apply(lambda x: x.strftime('%Y-%m-%d'))
        # data = data[data['t_warning'] == True].reset_index(drop=True)
        # data.loc[:, 't_warning'] = 1
        # data['c_code'] = data['c_code'].apply(bond_convert)
        tableName = 'b_stocka_warning'
        query_str = self.path + 'tableName=' + tableName + '&begDate=' + start_date + '&endDate=' + end_date
        if fields != None: query_str += '&fields=c_code,t_date,' + ','.join(fields)
        data = pd.read_csv(query_str)
        return data


# 获取基准指数
def get_index(start, end):
    data = w.wsd("889033.WI", "close", start, end, "Period=D;PriceAdj=B")
    df = pd.DataFrame(np.array(data.Data).T, index=data.Times, columns=['benchmark'])
    return df


# 提取股票相关数据#
def get_stock_data(stock_code, str_date):
    engine = create_engine('mysql+pymysql://hs_wangwenjie:hs_wangwenjie#A0@192.168.201.181:3306/wind?charset=utf8')
    # A股日行情表 S_DQ_VOLUME：成交量（一手100股）,S_DQ_AMOUNT成交额千元
    query1 = ("""select S_INFO_WINDCODE,TRADE_DT,S_DQ_CLOSE,S_DQ_HIGH,S_DQ_LOW,S_DQ_PCTCHANGE,S_DQ_VOLUME,S_DQ_TRADESTATUS,S_DQ_AMOUNT,S_DQ_ADJFACTOR
               from AShareEODPrices
               where S_INFO_WINDCODE in %s and TRADE_DT > '%s'""") % (stock_code, str_date)
    data1 = pd.read_sql_query(query1, engine).sort_values(by=['S_INFO_WINDCODE', 'TRADE_DT'])

    # A股日行情估值指标 S_VAL_MV总市值
    query2 = ("""select S_INFO_WINDCODE,TRADE_DT,S_VAL_MV,S_VAL_PE_TTM
               from AShareEODDerivativeIndicator
               where S_INFO_WINDCODE in %s and TRADE_DT > '%s'""") % (stock_code, str_date)

    data2 = pd.read_sql_query(query2, engine).sort_values(by=['S_INFO_WINDCODE', 'TRADE_DT'])
    # 获取正股行业代码 	SW_IND_CODE
    query3 = ("""select S_INFO_WINDCODE,SW_IND_CODE
               from AShareSWIndustriesClass
               where S_INFO_WINDCODE in %s""") % (stock_code,)
    data3 = pd.read_sql_query(query3, engine).sort_values(by=['S_INFO_WINDCODE'])
    # 获取正股ttm
    query4 = ("""select S_INFO_WINDCODE,TRADE_DT,S_DFA_NETPROFIT_TTM,S_DFA_DEDUCTEDPROFIT_TTM
               from PITFinancialFactor
               where S_INFO_WINDCODE in %s and TRADE_DT > '%s'""") % (stock_code, str_date)
    data4 = pd.read_sql_query(query4, engine).sort_values(by=['S_INFO_WINDCODE'])

    return data1, data2, data3, data4


# 提取行业相关数据#
def get_stock_industry():
    engine = create_engine('mysql+pymysql://hs_wangwenjie:hs_wangwenjie#A0@192.168.201.181:3306/wind?charset=utf8')
    # A股日行情表
    query1 = ("""select INDUSTRIESCODE,INDUSTRIESNAME
               from AShareIndustriesCode""")
    data1 = pd.read_sql_query(query1, engine)
    return data1


def get_stock_info(stock_code, str_date):
    engine = create_engine('mysql+pymysql://hs_wangwenjie:hs_wangwenjie#A0@192.168.201.181:3306/wind?charset=utf8')
    # A股利润表
    query1 = ("""select S_INFO_WINDCODE,ANN_DT,REPORT_PERIOD,OPER_REV,NET_PROFIT_EXCL_MIN_INT_INC,STATEMENT_TYPE
               from AShareIncome
               where S_INFO_WINDCODE in %s and ANN_DT > '%s'""") % (stock_code, str_date)
    data1 = pd.read_sql_query(query1, engine)

    # A股资产负债表
    query2 = ("""select S_INFO_WINDCODE,ANN_DT,TOT_SHRHLDR_EQY_EXCL_MIN_INT
               from AShareBalanceSheet
               where S_INFO_WINDCODE in %s and ANN_DT > '%s'""") % (stock_code, str_date)
    data2 = pd.read_sql_query(query2, engine)
    # 审计意见
    query3 = ("""select S_INFO_WINDCODE,ANN_DT,S_STMNOTE_AUDIT_CATEGORY
               from AShareAuditOpinion
               where S_INFO_WINDCODE in %s and ANN_DT > '%s' """) % (stock_code, str_date)
    data3 = pd.read_sql_query(query3, engine)

    return data1, data2, data3


# %% 提取数据
def get_data(method=1):
    # get_conv_data = ConvBondData()
    print('提取转债数据')
    path = os.getcwd()
    if method == 0:
        # df_v_s = pd.DataFrame()
        df_v_s_l = pd.DataFrame()
        df_s = pd.DataFrame()
        start_date = str(input("起始日期: "))
        # start_date = "2017-01-01"
    else:
        df_s = pd.read_csv("param/Mfactor_stock_data.csv")
        df_v_s_l = pd.read_csv("param/Mfactor_valuation_stock_last.csv")
        df_s.TRADE_DT = pd.to_datetime(df_s.TRADE_DT)
        df_v_s_l.Date = pd.to_datetime(df_v_s_l.Date) 
        start_date = str(max(df_v_s_l.Date))
        start_date = start_date[:10]

    # end_date = str(input("结束日期: "))
    end_date = str(date.today())  # 结束日期为今天
    # 提取债券价格数据、估值表 df_valuation
    df_valuation = ConvBondData().get_convbond_valuation(start_date, end_date)
    
    # 提取债券指数权重表 df_weight
    # df_weight = ConvBondData().get_index_weight(start_date, end_date)
    
    # 提取债券日行情表 df_market
    df_market = ConvBondData().get_convbond_market(start_date, end_date)
    
    # 提取债券信息 df_info
    # df_info = ConvBondData().get_convbond_info(start_date, end_date)
    
    df_warning = ConvBondData().get_stocka_warning(start_date, end_date)
    df_warning = df_warning.rename(columns={'t_tradingDate': 't_warning'})
    df_warning.t_warning = np.full([df_warning.shape[0]], 1)
    '''
    # 暂未更新股票数据导出
    
    # 提取股票日行情表 df_stock_market
    df_stock_market = ConvBondData().get_stocka_marketday(start_date, end_date)
    
    # 提取股票日预测指标表 df_stock_indicator
    df_stock_indicator = ConvBondData().get_stocka_derivativeindicator(start_date, end_date)
    
    # 提取股票行业 df_stock_info
    df_stock_info = ConvBondData().get_stocka_plateinfo(start_date, end_date)
    '''
    
    # dataframe合并
    # 将债券数据（df_valuation）与正股指标匹配
    df_valuation = df_valuation.rename(columns={'c_underlyingCode': 'c_code'})
    
    df_valuation_stock = pd.merge(df_valuation, df_market, on=['c_bondCode', 'index'], how='left')     # 合并可转债估值与可转债日行情值等信息
    # df_valuation_stock = pd.merge(df_valuation_stock, df_info, on = ['c_bondCode'], how = 'left')     # 合并可转债与可转债基本信息等信息
    df_valuation_stock = pd.merge(df_valuation_stock, df_warning, on=['c_code', 'index'], how='left')
    df_valuation_stock = df_valuation_stock.rename(columns={'index_x': 'index'})
    df_valuation_stock['index'] = pd.to_datetime(df_valuation_stock['index'])
    
    
    print('提取正股数据')
    # 提取正股数据，包括成交额、行业
    stockCode = df_valuation_stock['c_code'].unique()
    data = get_stock_data(tuple(stockCode), start_date[:4] + start_date[5:7] + start_date[8:])
    
    df_stock = pd.merge(data[0], data[1], on=['S_INFO_WINDCODE', 'TRADE_DT'], how='left')
    df_stock = pd.merge(df_stock, data[2], on=['S_INFO_WINDCODE'], how='left')
    df_stock = pd.merge(df_stock, data[3], on=['S_INFO_WINDCODE', 'TRADE_DT'], how='left')
    df_stock.TRADE_DT = pd.to_datetime(df_stock.TRADE_DT)
    df_stock_data = data[0]
    df_stock_data.TRADE_DT = pd.to_datetime(df_stock_data.TRADE_DT)
    # df_valuation_stock = pd.merge(df_valuation_stock, df_stock_indicator, on = ['c_code', 'index'], how = 'left')     # 合并可转债与正股估值等信息
    # df_valuation_stock = pd.merge(df_valuation_stock, df_stock_info, on = ['c_code', 'index'], how = 'left')     # 合并正股行业信息
    df_stock = df_stock.rename(columns={'S_INFO_WINDCODE': 'c_code', 'TRADE_DT': 'index'})
    df_valuation_stock = pd.merge(df_valuation_stock, df_stock, on=['c_code', 'index'], how='left')
    
    # 将日期列名改为Date
    df_valuation_stock = df_valuation_stock.rename(columns={'index': 'Date'})
    df_valuation_stock = df_valuation_stock.drop_duplicates(subset=("Date", 'c_bondCode'))
    df_valuation_stock = df_valuation_stock.sort_values(by="Date")

    # 获取月份最后一天的日期的代码
    last_day = df_valuation_stock["Date"].unique()
    df_valuation_stock_last = df_valuation_stock[df_valuation_stock['Date'].isin(last_day)]
    dt_ind = get_stock_industry()
    dt_ind = dt_ind.rename(columns={'INDUSTRIESCODE': 'SW_IND_CODE'})
    df_valuation_stock_last.SW_IND_CODE = df_valuation_stock_last.SW_IND_CODE.str[:4] + 12 * '0'
    df_valuation_stock_last = pd.merge(df_valuation_stock_last, dt_ind, on=['SW_IND_CODE'], how='left')
    industries = list(set(df_valuation_stock_last.INDUSTRIESNAME))[1:]
    date_n = len(last_day)

    # 提取股票信息
    data2 = get_stock_info(tuple(stockCode), "2017-01-01")
    data0 = data2[0][(data2[0]["STATEMENT_TYPE"] == "408001000") & (data2[0].REPORT_PERIOD.apply(lambda x: x[4:] == "1231"))]
    data0 = data0[data0.REPORT_PERIOD != np.nan]
    df_stock_info = pd.merge(data0, data2[1], on=['S_INFO_WINDCODE', 'ANN_DT'], how='left')
    df_stock_info = pd.merge(df_stock_info, data2[2], on=['S_INFO_WINDCODE', 'ANN_DT'], how='right')
    df_stock_info = df_stock_info.drop_duplicates(subset=("ANN_DT", 'S_INFO_WINDCODE'))
    df_stock_info.ANN_DT = pd.to_datetime(df_stock_info.ANN_DT)
    result = pd.DataFrame()

    print("源数据输出")
    # info2是需要输出的数据
    for j in tqdm(range(stockCode.shape[0])):
        temp_info = df_stock_info[df_stock_info.S_INFO_WINDCODE == stockCode[j]]
        temp_date = np.sort(temp_info.ANN_DT)
        temp_n = temp_info.shape[0]
        temp_res = df_valuation_stock[df_valuation_stock.c_code == stockCode[j]][["Date", "c_code"]]
        temp_res[df_stock_info.columns] = np.ones(shape=[temp_res.shape[0], len(df_stock_info.columns)])
    
        for k in range(1, temp_n):
            temp_len = np.sum((temp_res.Date < temp_date[k]) & (temp_res.Date >= temp_date[k - 1]))
            if temp_len == 0:
                continue
            temp = np.repeat(np.array(temp_info[temp_info.ANN_DT == temp_date[k - 1]]), temp_len, axis=0)
            temp_res.loc[(temp_res.Date < temp_date[k]) & (temp_res.Date >= temp_date[k - 1]), df_stock_info.columns] = list(temp)
            
        temp_len = np.sum(temp_res.Date >= temp_date[temp_n - 1])
        if temp_len == 0:
            result = pd.concat([result, temp_res])
            continue
        temp = np.repeat(np.array(temp_info[temp_info.ANN_DT == temp_date[temp_n - 1]]), temp_len, axis=0)
        # temp.Date = temp_res.Date[temp_res.Date >= temp_date[temp_n-1]]
        temp_res.loc[temp_res.Date >= temp_date[temp_n - 1], df_stock_info.columns] = list(temp)
        result = pd.concat([result, temp_res])
        
    result = result[result.iloc[:, -1] != 1]
    df_valuation_stock_last = pd.merge(df_valuation_stock_last, result, on=['Date', 'c_code'], how='left')
    # df_valuation_stock_last = df_valuation_stock_last.dropna()
    df_valuation_stock_last = df_valuation_stock_last.drop_duplicates(subset=("Date", 'c_bondCode'))
    
    df_stock_data = pd.concat([df_s, df_stock_data])
    df_valuation_stock_last = pd.concat([df_v_s_l, df_valuation_stock_last])
    # df_valuation_stock = pd.concat([df_v_s, df_valuation_stock])
    df_stock_data = df_stock_data.drop_duplicates(subset=("S_INFO_WINDCODE", 'TRADE_DT'))
    df_valuation_stock_last = df_valuation_stock_last.drop_duplicates(subset=("Date", 'c_bondCode'))
    # df_valuation_stock = df_valuation_stock.drop_duplicates(subset=("Date", 'c_code'))
    df_stock_data.to_csv(r"param/Mfactor_stock_data.csv", index=False)
    df_valuation_stock_last.to_csv(r"param/Mfactor_valuation_stock_last.csv", index=False)

    # 提取日期数据
    date_all = df_valuation_stock["Date"].unique()
    date_all = np.sort(date_all)
    return df_valuation_stock_last, df_stock_data, date_n, last_day, date_all
